#!/usr/bin/python

""" Get all data from Cassandra 0.6.8, then put them into cassandra 0.7.4
argv[1]: dst keySpace
"""

import os, sys, time
import pycassa
from thrift import Thrift  
from thrift.transport import TSocket  
from thrift.transport import TTransport  
from thrift.protocol import TBinaryProtocol
from cassandra import Cassandra
from cassandra.ttypes import *

if __name__ == '__main__':
    cassandra_6_srv_addr = '128.111.55.223'
    cassandra_6_port = 9160
    cassandra_7_srv_addr = '128.111.55.223'
    cassandra_7_port = 29160
    keyspace_name = "Keyspace1"
    column_family_name = "Standard1"
    dst_keyspace_name = sys.argv[1]
    
    # get data form cassandra 0.6
    print '====================================================================================='
    print 'Connecting to ' + cassandra_6_srv_addr + ':' + str(cassandra_6_port)
    transport = TTransport.TBufferedTransport(TSocket.TSocket(cassandra_6_srv_addr, cassandra_6_port))
    protocol = TBinaryProtocol.TBinaryProtocol(transport)
    client = Cassandra.Client(protocol)
    transport.open()
    print 'Connected!'
    
    print 'Connecting to ' + cassandra_7_srv_addr + ':' + str(cassandra_7_port)
    cass_7_pool = pycassa.connect(dst_keyspace_name, [cassandra_7_srv_addr + ':' + str(cassandra_7_port)])
    print 'Connected!'
    sys_manage = pycassa.system_manager.SystemManager(cassandra_7_srv_addr + ':' + str(cassandra_7_port))    

    # create keyspace_name if does not exist
    keyspace_list = sys_manage.list_keyspaces()
    if dst_keyspace_name not in keyspace_list:
        print 'Keyspace "' + dst_keyspace_name + '" does NOT exist, create it'
        sys_manage.create_keyspace(dst_keyspace_name, 1)

    # create a new column_family_dict 
    column_family_dict = sys_manage.get_keyspace_column_families(dst_keyspace_name)
    if column_family_dict.has_key(column_family_name):
        print 'Drop column family "' + column_family_name + '"'
        sys_manage.drop_column_family(dst_keyspace_name, column_family_name)
    print 'Create column family "' + column_family_name + '"'
    sys_manage.create_column_family(dst_keyspace_name, column_family_name)

    cass_7_col_fam = pycassa.ColumnFamily(cass_7_pool, column_family_name)
    
    # list all keys from cassandra 0.6
    column_parent = ColumnParent(column_family=column_family_name)

    predicate = SlicePredicate(column_names=['Encoded_Entity', 'Txn_Num'])
    key_range = KeyRange("ggreetingtests___Greeting___/", "ggreetingtests___Greeting___/~")
    start_key = "ggreetingtests___Greeting___/"
    end_key = "ggreetingtests___Greeting___/~"
    keyslices = client.get_range_slice(
                keyspace_name,
                column_parent,
                predicate,
                start_key,
                end_key,
                100000000,
                ConsistencyLevel.QUORUM) 
    print 'key number in 6: ' + str(len(keyslices))       
  
    count = 0
    for keyslice in keyslices:
        key_str = str(keyslice.key)
	
	for superColumn in keyslice.columns:
            # put this item into cassandra 0.7
            cass_7_col_fam.insert(key_str, {superColumn.column.name : superColumn.column.value})
            count += 1
    
    print 'insert number in 7: ' + str(count)
    
    print 'Done!'
    
    
    
    
    
